#include "buffer.h"
#include "common.h"
#include <time.h>

int main() {
  aclInit(NULL);
  int ret = buffer_init(DATA_SIZE);
  if (ret != 0) {
    fprintf(stderr, "buffer_init failed\n");
    return -1;
  }
  aclrtStream stream;
  ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
  ACL_CHECK(aclrtCreateStream(&stream));

  aclrtEvent startEvent, endEvent;
  ACL_CHECK(aclrtCreateEventWithFlag(&startEvent, ACL_EVENT_TIME_LINE));
  ACL_CHECK(aclrtCreateEventWithFlag(&endEvent, ACL_EVENT_TIME_LINE));

  int32_t canAccessPeer;
  ACL_CHECK(aclrtDeviceCanAccessPeer(&canAccessPeer, FROM_DEVICE, TO_DEVICE));
  if (canAccessPeer == 0) {
    fprintf(stderr, "Devices cannot access each other\n");
    return -1;
  }

  ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
  ACL_CHECK(aclrtDeviceEnablePeerAccess(TO_DEVICE, 0));
  ACL_CHECK(aclrtSetDevice(TO_DEVICE));
  ACL_CHECK(aclrtDeviceEnablePeerAccess(FROM_DEVICE, 0));

  // Warm-up iterations
  for (int i = 0; i < WARMUP_ITERATIONS; i++) {
    ACL_CHECK(aclrtSetDevice(TO_DEVICE));
    ACL_CHECK(aclrtMemset(dstbuf_d, DATA_SIZE, 0, DATA_SIZE));
    ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
    ACL_CHECK(aclrtMemcpyAsync(dstbuf_d, DATA_SIZE, srcbuf_d, DATA_SIZE,
                               ACL_MEMCPY_DEVICE_TO_DEVICE, stream));
    ACL_CHECK(aclrtSynchronizeStream(stream));
    if (buffer_verify() != 0) {
      fprintf(stderr, "data verify failed\n");
      return -1;
    }
  }

  // Measured iterations
  float totalElapsedTime = 0.0f;
  double totalWallTime = 0.0;
  for (int i = 0; i < MEASURED_ITERATIONS; i++) {
    float elapsedTime = 0.0f;
    struct timespec startWall, endWall;
    ACL_CHECK(aclrtSetDevice(TO_DEVICE));
    ACL_CHECK(aclrtMemset(dstbuf_d, DATA_SIZE, 0, DATA_SIZE));
    ACL_CHECK(aclrtSetDevice(FROM_DEVICE));    clock_gettime(CLOCK_MONOTONIC, &startWall);
    ACL_CHECK(aclrtRecordEvent(startEvent, stream));
    ACL_CHECK(aclrtMemcpyAsync(dstbuf_d, DATA_SIZE, srcbuf_d, DATA_SIZE,
                               ACL_MEMCPY_DEVICE_TO_DEVICE, stream));

    ACL_CHECK(aclrtRecordEvent(endEvent, stream));
    ACL_CHECK(aclrtSynchronizeStream(stream));
    clock_gettime(CLOCK_MONOTONIC, &endWall);
    ACL_CHECK(aclrtEventElapsedTime(&elapsedTime, startEvent, endEvent));

    double wallTime = (endWall.tv_sec - startWall.tv_sec) * 1000.0 +
                      (endWall.tv_nsec - startWall.tv_nsec) / 1e6;
    totalWallTime += wallTime;
    if (buffer_verify() != 0) {
      fprintf(stderr, "data verify failed\n");
      return -1;
    }
    totalElapsedTime += elapsedTime;
  }

  totalElapsedTime /= MEASURED_ITERATIONS;
  totalWallTime /= MEASURED_ITERATIONS;

  printf("Average P2P Time: %.3f ms, Bandwidth: %.3f GB/s\n",
         totalElapsedTime,
         ((float)DATA_SIZE) / (totalElapsedTime * 1e6f));
  printf("Wall Clock Time: %.3f ms, Bandwidth: %.3f GB/s\n", totalWallTime,
         ((float)DATA_SIZE) / (totalWallTime * 1e6));

  ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
  ACL_CHECK(aclrtDestroyEvent(startEvent));
  ACL_CHECK(aclrtDestroyEvent(endEvent));
  ACL_CHECK(aclrtDestroyStream(stream));
  ret = buffer_destroy();
  if (ret != 0) {
    fprintf(stderr, "buffer_destroy failed\n");
    return -1;
  }
  aclFinalize();
  return 0;
}